What does this PR do ?

main branch: #3076
NOTE: The kernel will be migrated to TE in future updates, and the main branch will depend on this version of the kernel.

1. Fix: incorrect parameter reference in Linear-Cross-Entropy Fusion for non-shared embeddings

When using non-shared embeddings, Linear-Cross-Entropy Fusion incorrectly references the shared embedding weight. In this scenario, the return value of shared_embedding_or_output_weight should not be used. Instead, it should explicitly reference output_layer.weight.

Additionally, since Linear-Cross-Entropy Fusion does not invoke the output_layer module, m-core DDP’s parameter gathering check may fail. To address this, I subclassed output_layer (i.e., ColumnParallelLinear) and moved the Linear-Cross-Entropy logic into its forward method.

2. Improvement: Use FP32 accumulation in multi-chunk computation for better convergence

The previous implementation of Linear-Cross-Entropy Fusion split computation along the vocabulary dimension and accumulated gradients in BF16, using the operation:

d_hidden += d_logits_slice * weight_slice

This caused d_hidden to accumulate in BF16, which could lead to reduced numerical accuracy.

This PR fixes the issue by allocating an FP32 buffer for d_hidden during accumulation, then downcasting it back to BF16 after accumulation completes. This improves numerical stability and overall convergence.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.